Skip to content

[SPARK-58093][SQL] Push down LIMIT through UNION ALL to Data Source V2 scan#57200

Open
j1wonpark wants to merge 1 commit into
apache:masterfrom
j1wonpark:pushdown-limit-union
Open

[SPARK-58093][SQL] Push down LIMIT through UNION ALL to Data Source V2 scan#57200
j1wonpark wants to merge 1 commit into
apache:masterfrom
j1wonpark:pushdown-limit-union

Conversation

@j1wonpark

@j1wonpark j1wonpark commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

V2ScanRelationPushDown.pushDownLimit had no case for Union, so a LIMIT on top of a UNION ALL was never pushed to the underlying Data Source V2 scans through SupportsPushDownLimit; it fell through to case other => (other, false).

Note the plan shape: LimitPushDown already pushes a LocalLimit into each Union branch (operator level), but the source-level limit hint (pushDownLimit) stops at the Union.

This PR adds two cases to pushDownLimit:

  • case u: Union — recurse into each branch, pushing the limit as a partial limit so the limit operators above the union are kept.
  • case l @ LocalLimit(IntegerLiteral(_), _) — look through the per-branch LocalLimit that LimitPushDown inserts, so the limit reaches the scan.

Results are unchanged. UNION ALL does not de-duplicate rows, so the union of branches that each keep at most limit rows still contains the first limit rows of the whole union — the same reason the existing Union case in LimitPushDown is safe. The union may return more than limit rows, so the push is only partial and the limit operators stay in the plan.

Two things worth calling out for review:

  • A branch may fully push a top N and drop its Sort (the Sort case already returns a plan without it). This is fine because Union does not attach meaning to the order in which it emits its children, and the limit above the union still caps the result. LIMIT over an ORDER BY that sits above the union is not pushed at all, which a test covers.
  • Because the Union case always reports the limit as not removable, plans that combine LIMIT with OFFSET can call pushLimit on the same builder twice with the same value. This already happens today for partially pushed scans and is idempotent for the sources in the repository.

Why are the changes needed?

SELECT ... FROM a UNION ALL SELECT ... FROM b LIMIT n never delivers the LIMIT to the scans, so each source does more work than necessary. Any Data Source V2 connector implementing SupportsPushDownLimit benefits from receiving the limit per branch — e.g. for JDBC, each branch's pushed-down SQL then carries LIMIT n.

Does this PR introduce any user-facing change?

No. Query results are identical. The only observable difference is that the pushed limit now appears on each UNION ALL branch's scan (e.g. PushedLimit: LIMIT n in the plan, or a LIMIT clause in the JDBC SQL), reducing work at the source.

How was this patch tested?

Added five tests to JDBCV2Suite, each asserting the limit pushed to every branch scan:

  • a two-branch union, filtered to a single matching employee so the result is deterministic and can be checked with checkAnswer;
  • a three-branch union;
  • a union with a sorted branch, where the branch limit is pushed as a top N and the branch Sort is dropped;
  • a union combined with OFFSET, covering both LIMIT n OFFSET m (branches keep n + m rows) and limit(n).offset(m) (branches keep n rows);
  • a negative test: a LIMIT over an ORDER BY on top of the union is not pushed to the scans.

Ran the full JDBCV2Suite (86 tests) — no regression.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus)

…2 scan

`V2ScanRelationPushDown.pushDownLimit` only handled a scan directly under the
limit (optionally through `Project`/`Sort`), so a `LIMIT` on top of a
`UNION ALL` fell through to `case other` and was never delivered to the
underlying Data Source V2 scans via `SupportsPushDownLimit`. `LimitPushDown`
already pushes a `LocalLimit` into each union branch, so the operator-level
limit reaches the branches while the source-level limit hint does not.

Add `Union` and `LocalLimit` cases to `pushDownLimit` that recurse into each
branch and push the limit as a partial limit, always keeping the outer limit
operators. UNION ALL never de-duplicates rows, so the union of branches each
capped at `limit` still contains the first `limit` rows of the whole union.

Signed-off-by: Jiwon Park <jpark92@outlook.kr>
@j1wonpark j1wonpark force-pushed the pushdown-limit-union branch from d2b65ba to c9e8a97 Compare July 12, 2026 12:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant